{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D Object Detection using TAO BEVFusion\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## What is BEVFusion?\n",
    "\n",
    "[BEVFusion](https://arxiv.org/abs/2205.13542) is a state of the art 3D Object Detection model using Camera-LiDAR Fusion. To fuse 3D features from LiDAR sensor with features from 2D image, BEVFusion uses BEV Feature space as commond feature space from both sensors. BEVFusion proposed to use Camera-To-BEV View-Transform to estimate BEV Feature from given 2D image and fuse BEV Features from Camera and LiDAR with ConvFuser module. Lastly, it has 3D Detection head to perfrom 3D object detection using fused feature.\n",
    "\n",
    "### Sample prediction of TAO BEVFusion model\n",
    "<img align=\"center\" src=\"./sample.jpg\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Convert Kitti dataset into TAO BEVFusion compatible input format by parsing pedestrian class only. (KittiPerson Dataset)\n",
    "* Take a pretrained model and finetune an BEVFusion model on KittiPerson dataset\n",
    "* Evaluate the trained model\n",
    "* Run inference with the trained model and visualize the result\n",
    "\n",
    "At the end of this notebook, you will have generated a trained `bevfusion` model.\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of BEVFusion using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate a trained model](#head-5)\n",
    "6. [Visualize inferences](#head-6)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/bevfusion/results`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"bevfusion\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"bevfusion\", \"results\")\n",
    "os.environ[\"LOCAL_CACHE_DIR\"] = os.path.join(os.environ['HOME'], \".cache\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/bevfusion\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.|\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_CACHE_DIR\"],\n",
    "           \"destination\": \"/.cache\"\n",
    "       },\n",
    "   ],\n",
    "   \"Envs\": [\n",
    "        {\n",
    "            \"variable\": \"TAO_TOOLKIT_CACHE\",\n",
    "            \"value\": \"/.cache\",\n",
    "        }\n",
    "    ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info\n",
    "!tao model -h"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the KITTI detection dataset for the tutorial. To find more details please visit\n",
    "http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download the KITTI detection images (http://www.cvlibs.net/download.php?file=data_object_image_2.zip), labels(http://www.cvlibs.net/download.php?file=data_object_label_2.zip), velodyne LIDAR pointcloud(http://www.cvlibs.net/download.php?file=data_object_velodyne.zip) and LIDAR calibration file(http://www.cvlibs.net/download.php?file=data_object_calib.zip) to $DATA_DOWNLOAD_DIR.\n",
    " \n",
    "The data will then be extracted to have below structure.\n",
    "\n",
    "```bash\n",
    "│── ImageSets\n",
    "│── training\n",
    "│   ├──calib & velodyne & label_2 & image_2\n",
    "│── testing\n",
    "│   ├──calib & velodyne & image_2\n",
    "```\n",
    "\n",
    "\n",
    "The `testing` directory will not be used in this notebook as it has no labels. For the `training` dataset, we will convert the labeles into pkl format in following section.\n",
    "\n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download and Verify dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have gotten the download links in your email, please populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL`,  `KITTI_LABELS_DOWNLOAD_URL`, `KITTI_LIDAR_DOWNLOAD_DIR` and `KITTI_CALIB_DOWNLOAD_DIR`. This next cell, will download the data and place in `$HOST_DATA_DIR`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "\n",
    "os.environ[\"URL_IMAGES\"]=\"KITTI_IMAGES_DOWNLOAD_URL\"\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $HOST_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "\n",
    "os.environ[\"URL_LABELS\"]=\"KITTI_LABELS_DOWNLOAD_URL\"\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $HOST_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi\n",
    "\n",
    "os.environ[\"URL_LIDAR\"]=\"KITTI_LIDAR_DOWNLOAD_URL\"\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_velodyne.zip ]; then wget $URL_LIDAR -O $HOST_DATA_DIR/data_object_velodyne.zip; else echo \"velodyne archive already downloaded\"; fi \n",
    "\n",
    "os.environ[\"URL_CALIB\"]=\"CALIB_DOWNLOAD_URL\"\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_calib.zip ]; then wget $URL_CALIB -O $HOST_DATA_DIR/data_object_calib.zip; else echo \"calib archive already downloaded\"; fi \n",
    "\n",
    "# Download train/val split\n",
    "!mkdir -p $HOST_DATA_DIR/ImageSets\n",
    "!if [ ! -f $HOST_DATA_DIR//ImageSets/test.txt ]; then wget https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/test.txt --no-check-certificate --content-disposition -O $HOST_DATA_DIR//ImageSets/test.txt; else echo \"test.txt archive already downloaded\"; fi \n",
    "!if [ ! -f $HOST_DATA_DIR//ImageSets/train.txt ]; then wget https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/train.txt --no-check-certificate --content-disposition -O $HOST_DATA_DIR/ImageSets/train.txt; else echo \"train.txt archive already downloaded\"; fi \n",
    "!if [ ! -f $HOST_DATA_DIR//ImageSets/val.txt ]; then wget https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/val.txt --no-check-certificate --content-disposition -O $HOST_DATA_DIR/ImageSets/val.txt; else echo \"val.txt archive already downloaded\"; fi \n",
    "!if [ ! -f $HOST_DATA_DIR//ImageSets/trainval.txt ]; then wget https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/trainval.txt --no-check-certificate --content-disposition -O $HOST_DATA_DIR/ImageSets/trainval.txt; else echo \"trainval.txt archive already downloaded\"; fi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Verification of the dataset\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_velodyne.zip ]; then echo 'Velodyne zip file not found, please download.'; else echo 'Found Velodyne zip file.';fi\n",
    "!if [ ! -f $HOST_DATA_DIR/data_object_calib.zip ]; then echo 'Calib zip file not found, please download.'; else echo 'Found Calib zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Unpack \n",
    "!unzip -u $HOST_DATA_DIR/data_object_image_2.zip -d $HOST_DATA_DIR\n",
    "!unzip -u $HOST_DATA_DIR/data_object_label_2.zip -d $HOST_DATA_DIR\n",
    "!unzip -u $HOST_DATA_DIR/data_object_velodyne.zip -d $HOST_DATA_DIR\n",
    "!unzip -u $HOST_DATA_DIR/data_object_calib.zip -d $HOST_DATA_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Convert datset to required format"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Converting downloaded Kitti labels into TAOBEVFusion compatible pkl format. This convert will only parse pedestrian class from public Kitti dataset and zero-pad rotations for x, z coordinates as only y-rotation is provided in Kitti. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!tao model bevfusion convert -e $SPECS_DIR/convert.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Download pre-trained model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/bevfusion:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/bevfusion:1.0 --dest $LOCAL_PROJECT_DIR/bevfusion/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_PROJECT_DIR/bevfusion/bevfusion_v1.0/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * type: KittiPersonDataset\n",
    "    * root_dir: root path for dataset\n",
    "    * gt_box_type: label boundinb box coordinate. Options to choose lidar or camera\n",
    "    * default_cam_key: default camera name. Defaults to CAM2 in Kitti\n",
    "    * train_dataset: \n",
    "        * repeat_time: number of repeat for training dataset. Defaults to 2 for Kitti\n",
    "        * ann_file: path to the training annotation pkl file generated via bevfusion convert\n",
    "        * data_prefix: \n",
    "            * pts: directory prefix name for lidar point cloud data\n",
    "            * img: directory prefix name for camera image data\n",
    "        * batch_size: batch size for train dataloader\n",
    "        * workers: number of workers to do train data loading\n",
    "    * val_dataset:\n",
    "        * ann_file: path to the validation annotation pkl file generated via bevfusion convert\n",
    "        * data_prefix: \n",
    "            * pts: directory prefix name for lidar point cloud data\n",
    "            * img: directory prefix name for camera image data\n",
    "        * batch_size: batch size for val dataloader\n",
    "        * workers: number of workers to do val data loading\n",
    "* model: configure the model setting\n",
    "    * type: model name. Currenlty only supporting BEVFusion\n",
    "    * point_cloud_range: point cloud range for the data\n",
    "    * voxel_size: voxel size\n",
    "    * grid_size: grid size for detection head\n",
    "* train: configure the training hyperparameter\n",
    "    * num_gpus: number of gpus\n",
    "    * num_nodes: number of nodes (num_nodes=1 for single node)\n",
    "    * validation_interval: validation interval\n",
    "    * pretrained_checkpoint: pretrained checkpoint for finetuning\n",
    "    * max_epoch: number of epochs\n",
    "    * optimizer:\n",
    "        * type: optimizer name. Defaults to AdamW\n",
    "        * lr: learning rate for the model\n",
    "    * lr_scheduler: dictionary for learning rate scheduler\n",
    "\n",
    "Please refer to the TAO documentation about BEVFusion to get all the parameters that are configurable.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* Evaluation uses Kitti3D metrics. For more info, please refer to: https://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d\n",
    "\n",
    "* Unlike the [original BEVFusion paper](https://arxiv.org/abs/2205.13542), we used three rotation angles for the labels. As a result, our detection head learns three rotation anlges represenation. Therefore, our model is not compatible with the original model code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!tao model bevfusion train -e $SPECS_DIR/experiment.yaml \\\n",
    "                            train.pretrained_checkpoint=/workspace/tao-experiments/bevfusion/bevfusion_v1.0/tao3d_bevfusion_epoch4.pth \\\n",
    "                            results_dir=$RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('Trained checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "%env NUM_EPOCH=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "%env CHECKPOINT={tmp[0]}\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/bevfusion_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/bevfusion_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate a trained model <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this section, we run the `evaluate` tool to evaluate the trained model and produce the AP40 metric.\n",
    "\n",
    "We provide evaluate.yaml specification files to configure the evaluate parameters including:\n",
    "\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * type: KittiPersonDataset\n",
    "    * root_dir: root path for dataset\n",
    "    * gt_box_type: label boundinb box coordinate. Options to choose lidar or camera\n",
    "    * default_cam_key: default camera name. Defaults to CAM2 in Kitti\n",
    "    * test_dataset:\n",
    "        * ann_file: path to the test annotation pkl file generated via bevfusion convert\n",
    "        * data_prefix: \n",
    "            * pts: directory prefix name for lidar point cloud data\n",
    "            * img: directory prefix name for camera image data\n",
    "        * batch_size: batch size for evaluate dataloader\n",
    "        * workers: number of workers to do evaluate data loading\n",
    "* model: configure the model setting\n",
    "    * type: model name. Currenlty only supporting BEVFusion\n",
    "    * point_cloud_range: point cloud range for the data\n",
    "    * voxel_size: voxel size\n",
    "    * grid_size: grid size for detection head\n",
    "* evaluate:\n",
    "  * num_gpus: number of gpus to use for evaluate\n",
    "\n",
    "* **NOTE: We reported 3D_AP40_moderate_strict as our AP40. Please look for this metric in the evaluation results.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!tao model bevfusion evaluate \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/bevfusion_model.pth \\\n",
    "            results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Visualize Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models and visualize the results. The `inference` tool produces annotated image outputs and json files that contain prediction information.\n",
    "\n",
    "We provide inference.yaml specification files to configure the evaluate parameters including:\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * type: KittiPersonDataset\n",
    "    * root_dir: root path for dataset\n",
    "    * gt_box_type: label boundinb box coordinate. Options to choose lidar or camera\n",
    "    * default_cam_key: default camera name. Defaults to CAM2 in Kitti\n",
    "    * test_dataset:\n",
    "        * ann_file: path to the test annotation pkl file generated via bevfusion convert\n",
    "        * data_prefix: \n",
    "            * pts: directory prefix name for lidar point cloud data\n",
    "            * img: directory prefix name for camera image data\n",
    "        * batch_size: batch size for evaluate dataloader\n",
    "        * workers: number of workers to do evaluate data loading\n",
    "* model: configure the model setting\n",
    "    * type: model name. Currenlty only supporting BEVFusion\n",
    "    * point_cloud_range: point cloud range for the data\n",
    "    * voxel_size: voxel size\n",
    "    * grid_size: grid size for detection head\n",
    "* inference:\n",
    "    * num_gpus: number of gpus to use for evaluate\n",
    "    * conf_threshold: confidence scrore threshold to filter out low score predictions\n",
    "\n",
    "For running inference with single image and point cloud file, you can run with inference_single.yaml file. Note that cam2img and lidar2cam need to be provided as shown in the yaml file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!tao model bevfusion inference \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "            inference.checkpoint=$RESULTS_DIR/train/bevfusion_model.pth \\\n",
    "            results_dir=$RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    " !tao model bevfusion inference \\\n",
    "            -e $SPECS_DIR/inference_single.yaml \\\n",
    "            inference.checkpoint=$RESULTS_DIR/train/bevfusion_model.pth \\\n",
    "            results_dir=$RESULTS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
